#include "bsdf.h"
#include "rng.h"

Spectrum BSDF::f(const Vector3 &woW, const Vector3 &wiW, BxDFType flags) const {
	
	Spectrum f = 0.f;
	
	Vector3 wi = WorldToLocal(wiW);
	Vector3 wo = WorldToLocal(woW);

	// Handle light leak / dark spot well...
	// Due to shading normal...

	if(Dot(wi, ng) * Dot(wo, ng) >0 ) // ignore BTDFs
		flags = BxDFType(flags & ~BSDF_TRANSMISSION);
	else	// ignore BRDFs
		flags = BxDFType(flags & ~BSDF_REFLECTION);

	for(int i = 0; i < nBxDFs; i++) {
		if(bxdfs[i]->MatchesFlags(flags))
			f += bxdfs[i]->f(wo, wi);
	}

	return f;
}

Spectrum BSDF::rho(const Vector3 &woW, BxDFType flags /* = BSDF_ALL */) const {
	Spectrum f = 0.f;
	Vector3 wo = WorldToLocal(woW);
	for(int i = 0; i < nBxDFs; i++)
		if(bxdfs[i]->MatchesFlags(flags))
			f += bxdfs[i]->rho(wo);
	return f;
}

Spectrum BSDF::rho(BxDFType flags) const {
	Spectrum f = 0.f;
	for(int i = 0; i < nBxDFs; i++)
		if(bxdfs[i]->MatchesFlags(flags))
			f += bxdfs[i]->rho();
	return f;
}

Spectrum BSDF::Sample_f(const Vector3 &woW, Vector3 *wiW, BxDFType flags /* = BSDF_ALL */, BxDFType *sampledType /* = NULL */) const {
	RNG rng(time(NULL));
	float pdf;
	Point2 sample(rng.RandomFloat(), rng.RandomFloat());
	Spectrum f = Sample_f(woW, wiW, rng.RandomFloat(), sample, &pdf, flags, sampledType);
	if (!f.Black() && pdf > 0.) f /= pdf;
	return f;
}

Spectrum BSDF::Sample_f(const Vector3 &woW, Vector3 *wiW, const float u, const Point2 &sample, float *pdf, BxDFType flags, BxDFType *sampledType) const {
	// Choose which bxdf to sample
	int matchingComps = NumComponent(flags);
	if(matchingComps == 0) {
		*pdf = 0.f;
		return Spectrum(0.f);
	}
	int which = min(Floor2Int(u*matchingComps), matchingComps-1);
	BxDF *bxdf = NULL;
	int count = which;
	for(int i = 0; i < nBxDFs; i++)
		if(bxdfs[i]->MatchesFlags(flags))
			if(--count == 0) {
				bxdf = bxdfs[i];
				break;
			}
	// Sample chosen bxdf
	Vector3 wi, wo = WorldToLocal(woW);
	*pdf = 0.f;
	Spectrum f = bxdf->Sample_f(wo, &wi, sample, pdf);
	if(f == 0.f)
		return 0.f;
	if(sampledType) *sampledType = bxdf->type;
	*wiW = LocalToWorld(wi);
	// Compute overall PDF with all matched bxdfs
	if(!(bxdf->type & BSDF_SPECULAR) && matchingComps > 1) {
		for(int i = 0; i < nBxDFs; i++) {
			if(bxdfs[i] != bxdf && bxdfs[i]->MatchesFlags(flags))
				*pdf += bxdfs[i]->Pdf(wo, wi);
		}
	}
	if(matchingComps > 1)
		*pdf /= matchingComps;
	// Compute value of BSDF for sampled direction
	if(!(bxdf->type & BSDF_SPECULAR)) {
		f = 0.f;
		if(Dot(*wiW, ng) * Dot(woW, ng) > 0) // ignore BTDFs
			flags = BxDFType(flags & ~BSDF_TRANSMISSION);
		else
			flags = BxDFType(flags & ~BSDF_REFLECTION);
		for(int i = 0; i < nBxDFs; i++)
			if(bxdfs[i]->MatchesFlags(flags))
				f += bxdfs[i]->Pdf(wo, wi);
	}

	return f;
}

float BSDF::Pdf(const Vector3 &woW, const Vector3 &wiW, BxDFType flags /* = BSDF_ALL */) const {
	if (nBxDFs == 0.) return 0.;
	Vector3 wo = WorldToLocal(woW), wi = WorldToLocal(wiW);
	float pdf = 0.f;
	int matchingComps = 0;
	for (int i = 0; i < nBxDFs; ++i)
		if (bxdfs[i]->MatchesFlags(flags)) {
			++matchingComps;
			pdf += bxdfs[i]->Pdf(wo, wi);
		}
	return matchingComps > 0 ? pdf / matchingComps : 0.f;
}

MemoryPool BSDF::memorypool;